#!/usr/bin/env python
# -*- coding: UTF-8 -*-
"""
@IDE     ：PyCharm 
@Author  ：kiway
@Date    ：2024/11/9 17:03 
"""

import asyncio
import traceback

import zmq.asyncio
import asyncio
import sys
from ..config import config
from ..logger import pretty_print

if sys.platform == "win32":
    asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())


class Subscriber:

    def __init__(self):
        # 创建一个异步上下文
        self.context = zmq.asyncio.Context()

        # 创建一个订阅者套接字
        self.socket = self.context.socket(zmq.SUB)
        self.socket.connect(f"tcp://127.0.0.1:{config.mq_port}")  # 连接到发布者的地址

        # 订阅所有消息
        self.socket.setsockopt_string(zmq.SUBSCRIBE, "")
        # pretty_print(f"subscriber {config.mq_port} 初试化...")

    async def get(self):
        # 异步接收消息
        message = await self.socket.recv_string()
        return message


class Publisher:
    # 创建一个ZMQ上下文和一个PUB套接字

    def __init__(self):
        self.context = zmq.Context()
        self.socket = self.context.socket(zmq.PUB)
        self.socket.bind(f'tcp://127.0.0.1:{config.mq_port}')  # 绑定到特定端口
        pretty_print(f"publisher {config.mq_port} 初试化...")

    def put(self, message):
        try:
            self.socket.send_string(str(message))

        except:
            traceback.print_exc()
